# utils.py
import os
import signal
import logging
import time
from pathlib import Path
from PIL import Image
import imagehash
import config  # 导入config模块
from functools import wraps

logger = logging.getLogger(__name__)


# === 日志配置 ===
class CustomLogFormatter(logging.Formatter):
    printed_messages = set()

    def formatTime(self, record, datefmt=None):
        ct = self.converter(record.created)
        t = time.strftime("%Y-%m-%d %H:%M:%S", ct)
        return f"{t},{int(record.msecs):03d}"

    def format(self, record):
        msg = super().format(record)
        original = record.getMessage()
        if "开始运行测试用例:" in original:
            if original in self.printed_messages:
                return ""
            self.printed_messages.add(original)
        return msg


class ColorStreamHandler(logging.StreamHandler):
    COLOR_MAP = {
        logging.DEBUG: "\033[36m",  # 青色
        logging.INFO: "\033[37m",  # 白色
        logging.WARNING: "\033[33m",  # 黄色
        logging.ERROR: "\033[31m",  # 红色
        logging.CRITICAL: "\033[41m",  # 红底
        21: "\033[32m"  # 绿色（SUCCESS）
    }

    def emit(self, record):
        try:
            msg = self.format(record)
            color = self.COLOR_MAP.get(record.levelno, "\033[37m")
            self.stream.write(f"{color}{msg}\033[0m\n")
            self.flush()
        except Exception as e:
            print(f"日志输出异常: {str(e)}")


def setup_logging(log_file_path=None):
    """设置日志记录，添加控制台和文件处理器"""
    # 添加成功样式
    logging.addLevelName(21, "SUCCESS")
    logging.Logger.success = lambda inst, msg, *args, **kwargs: inst.log(21, msg, *args, **kwargs)

    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)

    # 清除现有处理器，避免重复日志
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)

    # 创建控制台处理器
    console_handler = ColorStreamHandler()
    console_handler.setLevel(logging.INFO)

    # 设置日志格式
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)

    # 如果提供了日志文件路径，添加文件处理器
    if log_file_path:
        # 确保日志目录存在
        log_file_path.parent.mkdir(parents=True, exist_ok=True)

        file_handler = logging.FileHandler(log_file_path, encoding='utf-8')
        file_handler.setLevel(logging.DEBUG)  # 文件记录所有级别的日志
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    else:
        # 使用默认日志路径（原配置）
        log_file = config.REPORT_DIR / "test_run.log"
        os.makedirs(config.REPORT_DIR, exist_ok=True)
        file_handler = logging.FileHandler(str(log_file), encoding="utf-8")
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

    # 屏蔽Airtest冗余日志
    logging.getLogger("airtest").setLevel(logging.WARNING)
    logging.getLogger("airtest.report").setLevel(logging.WARNING)


# === 中断处理 ===
def handle_interrupt(sig, frame):
    logger = logging.getLogger()
    logger.warning(f"检测到停止信号（信号编号：{sig}），正在结束当前用例并生成汇总报告...")
    config.INTERRUPTED = True  # 标记中断状态


# === 环境清理 ===
def perform_cleanup(dev):  # 添加 dev 参数
    """清理测试环境并验证主页一致性（使用已连接的设备对象）"""
    try:
        timestamp = time.strftime("%Y%m%d%H%M%S")

        # 直接使用传入的 dev 对象清理主应用
        dev.stop_app("cn.migu.video.clip.walle")
        time.sleep(2)

        # 保存清理后截图（使用 dev 的 snapshot 方法）
        cleaned_shot = config.CLEANUP_DIR / f"cleaned_home_{timestamp}.png"
        dev.snapshot(str(cleaned_shot))
        logging.getLogger().info(f"清理后主页截图：{cleaned_shot}")

        # 验证与基准图一致
        if compare_images(config.BASELINE_PATH, cleaned_shot):
            logging.getLogger().success("环境恢复成功（主页一致）")
            return True

        # 备选清理（相册）
        dev.stop_app("com.huawei.photos")
        time.sleep(2)
        after_album_clean = config.CLEANUP_DIR / f"after_album_clean_{timestamp}.png"
        dev.snapshot(str(after_album_clean))
        if compare_images(config.BASELINE_PATH, after_album_clean):
            logging.getLogger().success("通过清理相册恢复主页")
            return True

        logging.getLogger().error("双重清理后仍未恢复主页")
        return False
    except Exception as e:
        logging.getLogger().error(f"清理异常: {str(e)}")
        return False


def compare_images(img1_path, img2_path, threshold=5):
    """对比两张图片是否相似（基于哈希值）"""
    try:
        hash1 = imagehash.average_hash(Image.open(img1_path))
        hash2 = imagehash.average_hash(Image.open(img2_path))
        return abs(hash1 - hash2) <= threshold
    except Exception as e:
        logging.getLogger().error(f"图片比对失败: {str(e)}")
        return False


def retry(times=3, delay=1, exceptions=(Exception,)):
    """重试装饰器"""

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            attempt = 0
            while attempt < times:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    attempt += 1
                    logger.warning(f"操作失败 ({func.__name__}), 尝试 {attempt}/{times}: {str(e)}")
                    if attempt < times:
                        time.sleep(delay)
            # 最后一次尝试后仍然失败
            logger.error(f"操作 {func.__name__} 重试 {times} 次后仍失败")
            raise  # 或者返回 None

        return wrapper

    return decorator